/* 
* Copyright (C) 2010, John Asmuth

 Licensed under the Apache License, Version 2.0 (the "License");
 you may not use this file except in compliance with the License.
 You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

 Unless required by applicable law or agreed to in writing, software
 distributed under the License is distributed on an "AS IS" BASIS,
 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
* 
*  $Revision$
*  $Date$
*  $Author$
*  $HeadURL$
* 
*/

package rlglue

import (
	"bytes"
	"errors"
	"fmt"
	"net"
	"os"
	"strconv"
	"strings"
	"time"
	"unsafe"
)

const (
	// RL-Glue needs to know what type of object is trying to connect.
	kExperimentConnection  = 1
	kAgentConnection       = 2
	kEnvironmentConnection = 3

	kAgentInit    = 4 // agent_* start by sending one of these values
	kAgentStart   = 5 // to the client to let it know what type of
	kAgentStep    = 6 // event to respond to
	kAgentEnd     = 7
	kAgentCleanup = 8
	kAgentMessage = 10

	kEnvInit    = 11
	kEnvStart   = 12
	kEnvStep    = 13
	kEnvCleanup = 14
	kEnvMessage = 19

	kRLInit         = 20
	kRLStart        = 21
	kRLStep         = 22
	kRLCleanup      = 23
	kRLReturn       = 24
	kRLNumSteps     = 25
	kRLNumEpisodes  = 26
	kRLEpisode      = 27
	kRLAgentMessage = 33
	kRLEnvMessage   = 34

	kRLTerm = 35

	kLocalHost    = "127.0.0.1"
	kDefaultPort  = 4096
	kRetryTimeout = 2

	kDefaultBufferSize = 4096
	kIntSize           = 4
	kDoubleSize        = 8
	kCharSize          = 1
)

func GetEnvHostPort() (host string, port int) {
	host, err := os.Getenverror("RLGLUE_HOST")
	if err != nil {
		host = "localhost"
	}
	portstr, err := os.Getenverror("RLGLUE_PORT")
	if err != nil {
		port = 4096
		return
	}
	port, err = strconv.Atoi(portstr)
	if err != nil {
		port = 4096
	}
	return
}

func reverse4(buffer [4]byte) [4]byte {
	for i := 0; i < 2; i++ {
		buffer[i], buffer[3-i] = buffer[3-i], buffer[i]
	}
	return buffer
}

func reverse8(buffer [8]byte) [8]byte {
	for i := 0; i < 4; i++ {
		buffer[i], buffer[7-i] = buffer[7-i], buffer[i]
	}
	return buffer
}

type network struct {
	conn       *net.TCPConn
	sendBuffer *bytes.Buffer
	recvBuffer *bytes.Buffer
}

func (nw *network) Connect(host string, port int) (err error) {
	raddrStr := fmt.Sprintf("%s:%d", host, port)
	raddr, err := net.ResolveTCPAddr("tcp", raddrStr)
	if err != nil {
		return
	}

	for nw.conn == nil || err != nil {
		nw.conn, err = net.DialTCP("tcp", nil, raddr)
		if err != nil {
			time.Sleep(kRetryTimeout * 1000000000)
		}
	}

	nw.sendBuffer = bytes.NewBuffer(make([]byte, kDefaultBufferSize)[0:0])
	nw.recvBuffer = bytes.NewBuffer(make([]byte, kDefaultBufferSize)[0:0])

	return
}

func (nw *network) Close() (err error) {
	err = nw.conn.Close()
	return
}

func (nw *network) Send() (err error) {
	_, err = nw.sendBuffer.WriteTo(nw.conn)
	return
}

func (nw *network) Recv(size int32) (bytesRead int32, err error) {
	if size == 0 {
		return
	}
	bytesToRead := make([]byte, size)
	bytesReadInt, err := nw.conn.Read(bytesToRead)
	bytesRead = int32(bytesReadInt)
	if err == nil {
		nw.recvBuffer.Write(bytesToRead)
	}
	return
}

func (nw *network) ClearSendBuffer() {
	nw.sendBuffer.Truncate(0)
}

func (nw *network) ClearRecvBuffer() {
	nw.recvBuffer.Truncate(0)
}

func (nw *network) FlipSendBuffer() {
	nw.ClearSendBuffer()
}

func (nw *network) FlipRecvBuffer() {
	nw.ClearRecvBuffer()
}

func (nw *network) GetInt() (res int32, err error) {
	if nw.recvBuffer.Len() < kIntSize {
		err = errors.New("Not enough data buffered to read int")
		return
	}
	var intBytes [kIntSize]byte
	nw.recvBuffer.Read(intBytes[0:kIntSize])
	intBytes = reverse4(intBytes)
	res = *(*int32)(unsafe.Pointer(&intBytes))
	return
}

func (nw *network) PutInt(value int32) {
	intBytes := *(*[kIntSize]byte)(unsafe.Pointer(&value))
	intBytes = reverse4(intBytes)
	nw.sendBuffer.Write(intBytes[0:kIntSize])
}

func (nw *network) GetDouble() (res float64, err error) {
	if nw.recvBuffer.Len() < kDoubleSize {
		err = errors.New("Not enough data buffered to read double")
		return
	}
	var doubleBytes [kDoubleSize]byte
	nw.recvBuffer.Read(doubleBytes[0:kDoubleSize])
	doubleBytes = reverse8(doubleBytes)
	res = *(*float64)(unsafe.Pointer(&doubleBytes))
	return
}

func (nw *network) PutDouble(value float64) {
	doubleBytes := *(*[kDoubleSize]byte)(unsafe.Pointer(&value))
	doubleBytes = reverse8(doubleBytes)
	nw.sendBuffer.Write(doubleBytes[0:kDoubleSize])
}

func (nw *network) GetString() (res string, err error) {
	length, err := nw.GetInt()
	if err != nil {
		return
	}
	if int32(nw.recvBuffer.Len()) < length {
		err = errors.New("Not enough data buffered to read string")
		return
	}
	stringBytes := make([]byte, length)
	nw.recvBuffer.Read(stringBytes)
	res = string(stringBytes)
	return
}

func (nw *network) PutString(value string) {
	nw.PutInt(int32(len(value)))
	nw.sendBuffer.ReadFrom(strings.NewReader(value))
}

func (nw *network) getAbstractType() (ati AbstractType, err error) {
	numInts, err := nw.GetInt()
	if err != nil {
		return
	}
	numDoubles, err := nw.GetInt()
	if err != nil {
		return
	}
	numChars, err := nw.GetInt()
	if err != nil {
		return
	}

	ints := make([]int32, numInts)
	for i := int32(0); i < numInts; i++ {
		ints[i], err = nw.GetInt()
		if err != nil {
			return
		}
	}
	doubles := make([]float64, numDoubles)
	for i := int32(0); i < numDoubles; i++ {
		doubles[i], err = nw.GetDouble()
		if err != nil {
			return
		}
	}
	if int32(nw.recvBuffer.Len()) < numChars {
		err = errors.New("Not enough data buffered to read chars")
		return
	}
	chars := make([]byte, numChars)
	nw.recvBuffer.Read(chars)

	ati = newAT(ints, doubles, chars)
	return
}

func (nw *network) putAbstractType(at AbstractType) {
	nw.PutInt(int32(len(at.Ints())))
	nw.PutInt(int32(len(at.Doubles())))
	nw.PutInt(int32(len(at.Chars())))
	for _, val := range at.Ints() {
		nw.PutInt(val)
	}
	for _, val := range at.Doubles() {
		nw.PutDouble(val)
	}
	nw.sendBuffer.Write(at.Chars())
}

func (nw *network) sizeOfAbstractType(at AbstractType) (size int32) {
	size = kIntSize * 3
	size += kIntSize * int32(len(at.Ints()))
	size += kDoubleSize * int32(len(at.Doubles()))
	size += kCharSize * int32(len(at.Chars()))
	return
}

func (nw *network) GetObservation() (obs Observation, err error) {
	return nw.getAbstractType()
}

func (nw *network) PutObservation(obs Observation) {
	nw.putAbstractType(obs)
}

func (nw *network) GetAction() (act Action, err error) {
	return nw.getAbstractType()
}

func (nw *network) PutAction(action Action) {
	nw.putAbstractType(action)
}

func (nw *network) PutRewardObservation(obs Observation, r float64, t bool) {
	if t {
		nw.PutInt(1)
	} else {
		nw.PutInt(0)
	}
	nw.PutDouble(r)
	nw.PutObservation(obs)
}

func (nw *network) SizeOfAction(action Action) int32 {
	return nw.sizeOfAbstractType(action)
}

func (nw *network) SizeOfObservation(observation Observation) int32 {
	return nw.sizeOfAbstractType(observation)
}

func (nw *network) SizeOfRewardObservation(obs Observation) (size int32) {
	size = kIntSize + kDoubleSize + nw.SizeOfObservation(obs)
	return
}
